package com.jerry.hm.datastructure.blockingqueue;

import java.util.Arrays;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 双锁实现阻塞队列
 *
 * @author Jerry
 * @since 2024/10/6 17:01
 */
public class DoubleLockBlockingQueue<E> implements BlockingQueue<E> {

    private final E[] array;
    private int head;
    private int tail;
    private AtomicInteger size = new AtomicInteger(0);

    private ReentrantLock offerLock = new ReentrantLock();
    private Condition notFull = offerLock.newCondition();

    private ReentrantLock pollLock = new ReentrantLock();
    private Condition notEmpty = pollLock.newCondition();


    public DoubleLockBlockingQueue(int capacity) {
        array = (E[]) new Object[capacity];
    }

    /**
     * 添加一个元素，如果队列满则阻塞
     *
     * @param e
     * @throws InterruptedException
     */
    @Override
    public void offer(E e) throws InterruptedException {
        // 添加元素前的数量
        int count;
        offerLock.lockInterruptibly();
        try {
            while (isFull()) {
                notFull.await();
            }
            array[tail++] = e;
            if (tail == array.length) {
                tail = 0;
            }
            count = size.getAndIncrement();
            if (count < array.length - 1) {
                notFull.signal();
            }
        } finally {
            offerLock.unlock();
        }

        // 唤醒等待非空的线程
        if (count == 0) {
            pollLock.lock();
            try {
                notEmpty.signal();
            } finally {
                pollLock.unlock();
            }
        }
    }

    /**
     * 在指定时间内添加一个元素，如果队列满则阻塞
     *
     * @param e
     * @param timeout 毫秒数
     * @return 添加成功或者失败
     * @throws InterruptedException
     */
    @Override
    public boolean offer(E e, long timeout) throws InterruptedException {
        // 添加元素前的数量
        int count;
        offerLock.lockInterruptibly();
        try {
            long remainder = TimeUnit.MILLISECONDS.toNanos(timeout);
            while (isFull()) {
                if (remainder <= 0) {
                    return false;
                }
                remainder = notFull.awaitNanos(remainder);
            }
            array[tail++] = e;
            if (tail == array.length) {
                tail = 0;
            }
            count = size.getAndIncrement();
            if (count < array.length - 1) {
                notFull.signal();
            }
        } finally {
            offerLock.unlock();
        }

        // 唤醒等待非空的线程
        if (count == 0) {
            pollLock.lock();
            try {
                notEmpty.signal();
            } finally {
                pollLock.unlock();
            }
        }
        return true;
    }

    /**
     * 移除一个元素，如果队列空则阻塞
     *
     * @return
     * @throws InterruptedException
     */
    @Override
    public E poll() throws InterruptedException {
        E e;
        // 移除前的元素个数
        int count;
        pollLock.lockInterruptibly();
        try {
            while (isEmpty()) {
                notEmpty.await();
            }
            e = array[head];
            array[head] = null;
            if (++head == array.length) {
                head = 0;
            }
            count = size.getAndDecrement();
            if (count > 1) {
                notEmpty.signal();
            }
        } finally {
            pollLock.unlock();
        }

        // 唤醒等待不满的线程
        if (count == array.length) {
            offerLock.lock();
            try {
                notFull.signal();
            } finally {
                offerLock.unlock();
            }
        }
        return e;
    }

    /**
     * 移除一个元素，如果队列空则阻塞
     *
     * @return
     * @throws InterruptedException
     */
    @Override
    public E poll(long timeout) throws InterruptedException {
        E e;
        // 移除前的元素个数
        int count;
        pollLock.lockInterruptibly();
        try {
            long remainder = TimeUnit.MILLISECONDS.toNanos(timeout);
            while (isEmpty()) {
                if (remainder <= 0) {
                    return null;
                }
                remainder = notEmpty.awaitNanos(remainder);
            }
            e = array[head];
            array[head] = null;
            if (++head == array.length) {
                head = 0;
            }
            count = size.getAndDecrement();
            if (count > 1) {
                notEmpty.signal();
            }
        } finally {
            pollLock.unlock();
        }

        // 唤醒等待不满的线程
        if (count == array.length) {
            offerLock.lock();
            try {
                notFull.signal();
            } finally {
                offerLock.unlock();
            }
        }
        return e;
    }

    private boolean isEmpty() {
        return size.get() == 0;
    }

    private boolean isFull() {
        return size.get() == array.length;
    }

    @Override
    public String toString() {
        return "array=" + Arrays.toString(array) + '}';
    }
}
